# -*- coding: utf-8 -*-
"""
基于LangChain的GraphRAG服务
使用阿里云百炼Qwen模型 + Neo4j知识图谱
"""

import os
import logging
from typing import Dict, Any, Optional, List
from langchain_community.graphs import Neo4jGraph
from langchain.chains import GraphCypherQAChain
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate

logger = logging.getLogger(__name__)


class LangChainGraphRAGService:
    """
    LangChain GraphRAG服务
    
    核心能力:
    1. 自然语言 → Cypher查询（自动生成）
    2. 图谱检索 → 结构化数据
    3. LLM增强 → 自然语言分析（支持Markdown格式）
    """
    
    def __init__(self):
        """初始化LangChain GraphRAG服务"""
        # 1. 连接Neo4j图数据库（APOC已安装）
        self.graph = Neo4jGraph(
            url=os.getenv('NEO4J_URI', 'bolt://43.140.197.100:7687'),
            username=os.getenv('NEO4J_USERNAME', 'neo4j'),
            password=os.getenv('NEO4J_PASSWORD', '12345678'),
            database=os.getenv('NEO4J_DATABASE', 'neo4j')
        )
        
        # 2. 初始化LLM（阿里云百炼Qwen）
        self.llm = ChatOpenAI(
            model="qwen3-max-2025-09-23",
            api_key=os.getenv("DASHSCOPE_API_KEY"),
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
            temperature=0,  # 保证输出稳定性
            streaming=False
        )
        
        # 3. 刷新图谱Schema（让LLM理解数据结构）
        try:
            self.graph.refresh_schema()
            logger.info(f"✅ Neo4j图谱Schema加载成功")
            logger.info(f"Schema: {self.graph.schema}")
        except Exception as e:
            logger.error(f"❌ 图谱Schema加载失败: {e}")
        
        # 4. 创建GraphCypherQAChain
        self._init_chain()
    
    def _init_chain(self):
        """初始化QA Chain，优化Prompt"""
        
        # Cypher生成Prompt（针对汽车销量数据优化）
        cypher_generation_template = """
你是一个Neo4j Cypher专家。基于Neo4j图谱Schema和用户问题，生成正确的Cypher查询。

Schema信息:
{schema}

重要说明:
1. **时间范围处理**: 
   - year_month字段格式为"YYYY-MM-DD"（如"2022-08-01"）
   - 使用 >= 和 < 进行范围查询
   - 例如：2022年 → WHERE sales.year_month >= "2022-01-01" AND sales.year_month < "2023-01-01"
   
2. **实体匹配**:
   - 车型名称直接使用，不要模糊匹配
   - 例如：海豹 → {{name: "海豹"}}

3. **销售数据查询**:
   - 销量字段：sales.units_sold
   - 价格字段：sales.low_price, sales.high_price
   - 时间字段：sales.year_month

4. **关系遍历**:
   - CarModel -[:HAS_SALES_RECORD]-> SalesRecord
   - CarModel -[:BELONGS_TO]-> Brand
   - Brand -[:OWNED_BY]-> Manufacturer
   - Brand -[:FROM_COUNTRY]-> Country
   - CarModel -[:HAS_BODY_TYPE]-> BodyType

5. **聚合统计**:
   - 总销量：sum(sales.units_sold)
   - 平均价格：avg(sales.low_price)
   - 时间排序：ORDER BY sales.year_month

用户问题: {question}

只返回Cypher查询，不要有任何解释。
Cypher查询:
"""
        
        cypher_prompt = PromptTemplate(
            input_variables=["schema", "question"],
            template=cypher_generation_template
        )
        
        # 回答生成Prompt（生成Markdown格式分析）
        qa_generation_template = """
你是一个专业的数据分析师。基于以下从Neo4j图谱检索的数据，生成专业的分析报告。

用户问题: {question}

图谱检索数据:
{context}

要求:
1. **使用Markdown格式**输出
2. **包含以下部分**:
   - ## 数据概览（用表格展示关键数据）
   - ## 趋势分析（描述变化趋势）
   - ## 关键洞察（3-5条要点）
   - ## 结论与建议（可选）

3. **数据展示**:
   - 使用Markdown表格格式化数据
   - 数字使用千分位逗号（如：15,378）
   - 增长率保留1位小数

4. **分析质量**:
   - 基于数据而非猜测
   - 如果数据不足，明确说明
   - 专业、简洁、可操作

开始分析:
"""
        
        qa_prompt = PromptTemplate(
            input_variables=["context", "question"],
            template=qa_generation_template
        )
        
        # 创建Chain
        self.chain = GraphCypherQAChain.from_llm(
            llm=self.llm,
            graph=self.graph,
            verbose=True,
            return_intermediate_steps=True,
            cypher_prompt=cypher_prompt,
            qa_prompt=qa_prompt,
            allow_dangerous_requests=True  # 允许执行查询
        )
        
        logger.info("✅ GraphCypherQAChain初始化成功")
    
    async def query(
        self,
        question: str,
        return_cypher: bool = True,
        return_context: bool = True
    ) -> Dict[str, Any]:
        """
        执行GraphRAG查询
        
        Args:
            question: 用户自然语言问题
            return_cypher: 是否返回生成的Cypher查询
            return_context: 是否返回图谱检索的原始数据
            
        Returns:
            {
                "answer": "Markdown格式的分析报告",
                "cypher": "生成的Cypher查询（可选）",
                "context": "图谱检索的原始数据（可选）",
                "metadata": {
                    "question": "原始问题",
                    "model": "使用的模型"
                }
            }
        """
        try:
            logger.info(f"📝 收到查询: {question}")
            
            # 执行Chain
            result = self.chain.invoke({"query": question})
            
            # 提取结果
            answer = result.get("result", "抱歉，无法生成分析。")
            
            # 构建响应
            response = {
                "answer": answer,
                "metadata": {
                    "question": question,
                    "model": "qwen3-max-2025-09-23",
                    "graph_database": "neo4j"
                }
            }
            
            # 添加中间步骤
            if return_cypher or return_context:
                intermediate_steps = result.get("intermediate_steps", [])
                
                if intermediate_steps:
                    # 提取Cypher查询
                    if return_cypher and len(intermediate_steps) > 0:
                        cypher_query = intermediate_steps[0].get("query", "")
                        response["cypher"] = cypher_query
                        logger.info(f"🔍 生成的Cypher: {cypher_query}")
                    
                    # 提取图谱上下文
                    if return_context and len(intermediate_steps) > 0:
                        context = intermediate_steps[0].get("context", [])
                        response["context"] = context
                        logger.info(f"📊 检索数据量: {len(context) if isinstance(context, list) else 'N/A'}")
            
            logger.info(f"✅ 查询完成")
            return response
            
        except Exception as e:
            logger.error(f"❌ GraphRAG查询失败: {e}")
            import traceback
            traceback.print_exc()
            
            return {
                "answer": f"查询失败: {str(e)}",
                "error": str(e),
                "metadata": {
                    "question": question,
                    "status": "error"
                }
            }
    
    async def get_graph_stats(self) -> Dict[str, Any]:
        """获取图谱统计信息"""
        try:
            # 总节点数
            total_nodes_query = "MATCH (n) RETURN count(n) as count"
            total_nodes = self.graph.query(total_nodes_query)[0]["count"]
            
            # 总关系数
            total_rels_query = "MATCH ()-[r]->() RETURN count(r) as count"
            total_rels = self.graph.query(total_rels_query)[0]["count"]
            
            # 节点类型统计
            node_labels_query = """
            CALL db.labels() YIELD label
            CALL {
                WITH label
                MATCH (n)
                WHERE label IN labels(n)
                RETURN count(n) as count
            }
            RETURN label, count
            ORDER BY count DESC
            """
            node_stats = self.graph.query(node_labels_query)
            
            # 关系类型统计
            rel_types_query = """
            CALL db.relationshipTypes() YIELD relationshipType
            CALL {
                WITH relationshipType
                MATCH ()-[r]->()
                WHERE type(r) = relationshipType
                RETURN count(r) as count
            }
            RETURN relationshipType, count
            ORDER BY count DESC
            """
            rel_stats = self.graph.query(rel_types_query)
            
            return {
                "total_nodes": total_nodes,
                "total_relationships": total_rels,
                "node_statistics": {item["label"]: item["count"] for item in node_stats},
                "relationship_statistics": {item["relationshipType"]: item["count"] for item in rel_stats},
                "schema": self.graph.schema
            }
            
        except Exception as e:
            logger.error(f"获取图谱统计失败: {e}")
            return {"error": str(e)}
    
    async def test_connection(self) -> Dict[str, Any]:
        """测试连接"""
        try:
            result = self.graph.query("RETURN 1 as test")
            return {
                "status": "connected",
                "test_result": result,
                "schema_loaded": bool(self.graph.schema)
            }
        except Exception as e:
            return {
                "status": "error",
                "error": str(e)
            }

